# -*- coding: utf-8 -*-
import numpy as np
import pandas as pd
import os
import pickle
import gradio as gr

class NaiveBayesClassifier:
    def __init__(self):
        self.priors = None
        self.conditionals = None

    # 训练朴素贝叶斯分类器
    def fit(self, X, y):
        num_classes = len(np.unique(y))  # 获取类别数
        num_features = X.shape[1]  # 获取特征数
        self.priors = np.zeros(num_classes)  # 初始化先验概率
        self.conditionals = np.zeros((num_classes, num_features))  # 初始化条件概率
        # 计算先验概率和条件概率
        for i in range(num_classes):
            X_i = X[y == i]
            self.priors[i] = X_i.shape[0] / X.shape[0]
            self.conditionals[i] = (X_i.sum(axis=0) + 1) / (X_i.sum() + num_features)

    # 预测分类结果
    def predict(self, X):
        log_probs = np.log(self.priors) + X @ np.log(self.conditionals.T)
        return np.argmax(log_probs, axis=1)

# 过滤姓名中的非法字符
def filter_name(name):
    name = ''.join(filter(lambda x: x.isalpha() or x.isalnum(), name))  # 只保留字母和数字
    return name


# 读取数据并进行预处理
train_data = pd.read_csv('C:/Users/yangkeyan/Desktop/bayes/NameData.txt', sep='\t', header=None, names=['name', 'gender'], encoding='utf-8')
test_data = pd.read_csv('C:/Users/yangkeyan/Desktop/bayes/NameData2.txt', sep='\t', header=None, names=['name', 'gender'], encoding='utf-8')
# 数据预处理
train_data['name'] = train_data['name'].apply(lambda x: list(x))
train_data['gender'] = train_data['gender'].apply(lambda x: 0 if x == '女' else 1)

test_data['name'] = test_data['name'].apply(lambda x: list(x))
test_data['gender'] = test_data['gender'].apply(lambda x: 0 if x == '女' else 1)

def build_word_vector(text, filter_words):
    text = list(text)
    word_vector = {word: 0 for word in filter_words}
    for word in text:
        if word in word_vector:
            word_vector[word] += 1
    return pd.Series(word_vector)


if os.path.exists('all_vectors.pkl'):
    with open('filter_words.txt', 'r', encoding='utf-8') as f:
        filter_words = sorted(set([line.strip() for line in f.readlines()]))  # 读取特征词
    with open('all_vectors.pkl', 'rb') as f:
        all_vectors = pickle.load(f)  # 读取特征向量
# 否则，生成特征向量和分类器，并保存到文件中
else:
    # 合并训练集和测试集数据，并统计每个汉字的出现频率
    train_name_list = train_data['name'].tolist()
    test_name_list = test_data['name'].tolist()

    name_list = train_name_list + test_name_list  # 合并训练集和测试集数据
    word_freq_dict = {}
    for name in name_list:
        name = list(name)
        for word in name:
            if word in word_freq_dict:
                word_freq_dict[word] += 1
            else:
                word_freq_dict[word] = 1

    word_freq_df = pd.DataFrame({'word': list(word_freq_dict.keys()), 'freq': list(word_freq_dict.values())})
    word_freq_df.sort_values(by='freq', ascending=False, inplace=True)

    # 选取出现频率大于等于 40 的汉字作为特征词
    filter_words = [' '] + word_freq_df[word_freq_df['freq'] >= 40]['word'].tolist()
    filter_words = sorted(set(filter_words))

    # 将每个姓名转换为特征向量
    all_vectors = pd.concat([train_data['name'], test_data['name']], axis=0)
    all_vectors = all_vectors.apply(lambda x: build_word_vector(x, filter_words)).fillna(0)

    # 将特征词和特征向量保存到文件中
    with open('filter_words.txt', 'w', encoding='utf-8') as f:
        f.write('\n'.join(filter_words))
    with open('all_vectors.pkl', 'wb') as f:
        pickle.dump(all_vectors, f)
train_vectors = all_vectors.iloc[:len(train_data), :]
test_vectors = all_vectors.iloc[len(train_data):, :]
# 如果分类器已经存在，则直接读取
if os.path.exists('model.pkl'):
    with open('model.pkl', 'rb') as f:
        clf = pickle.load(f)  # 读取分类器
# 否则，训练分类器，并保存到文件中
else:
    clf = NaiveBayesClassifier()
    clf.fit(train_vectors.values, train_data['gender'].values)  # 训练分类器
    with open('model.pkl', 'wb') as f:
        pickle.dump(clf, f)

# 定义函数，根据姓名预测性别
def predict_gender(name):
    vector = build_word_vector(name, filter_words).values.reshape(1, -1)  # 将姓名转换为特征向量
    gender = '女' if clf.predict(vector)[0] == 0 else '男'  # 预测性别
    return gender

# 定义函数，根据批量人名预测性别，每两个字或三个字为一个性别
def predict_gender_batch(names):
    names_list = names.split()  # 以空格号分隔姓名
    genders = []
    current_name = ''
    for name in names_list:
        current_name += name
        if len(current_name) >= 2:  # 当姓名长度达到2个字或以上时进行性别预测
            gender = predict_gender(current_name)
            genders.append(gender)
            current_name = ''  # 清空当前姓名
    return ' '.join(genders)  # 以空格号将性别连接成字符串返回

def count_gender_results(gender_results):
    genders_list = gender_results.split()  # 以空格分隔男女性别
    count = {'女': 0, '男': 0}
    for gender in genders_list:
        count[gender] += 1
    return count

def predict_gender_batch_from_file(file_path):
    with open(file_path, 'r', encoding='utf-8') as f:
        lines = f.readlines()

    total_names = len(lines)
    correct_count_male = 0
    correct_count_female = 0
    gender_results = ""

    male_count = 0
    female_count = 0
    predictions = []  # 存储每次预测的性别结果

    for line in lines:
        name, actual_gender = line.strip().split('\t')
        predicted_gender = predict_gender(name)
        gender_results += predicted_gender + " "  # 将预测的性别添加到结果中
        predictions.append(predicted_gender)  # 存储预测结果

        if predicted_gender == actual_gender:
            if actual_gender == '男':
                correct_count_male += 1
            elif actual_gender == '女':
                correct_count_female += 1

    # 计算准确率
    accuracy_male = correct_count_male / total_names * 100 if total_names > 0 else 0
    accuracy_female = correct_count_female / total_names * 100 if total_names > 0 else 0

    # 将预测结果写入当前路径下的genders.txt文件中
    with open('genders.txt', 'w', encoding='utf-8') as result_file:
        for prediction in predictions:
            result_file.write(f"{prediction}\n")
    
    print()

    # 准备结果摘要
    result_summary = f"预测结果已写入当前路径下的genders.txt文件中" \
                     f"女姓名数量: {gender_results.count('女')}, 男姓名数量: {gender_results.count('男')}, " \
                     f"女姓名正确数量: {correct_count_female}, 男姓名正确数量: {correct_count_male}, " \
                     f"女姓名正确率: {accuracy_female:.2f}%, 男姓名正确率: {accuracy_male:.2f}%"

    return result_summary

# 创建单个人名性别判定的 Gradio 接口
iface1 = gr.Interface(
    predict_gender,
    inputs="text",
    outputs="text",
    title="单个人名性别判定",
    description="输入单个人名进行性别判定"
)

# 添加批量人名性别判定的接口
iface2 = gr.Interface(
    predict_gender_batch,
    inputs="text",
    outputs="text",
    title="批量人名性别判定",
    description="输入批量人名进行性别判定，批量人名之间以空格为分隔符，可跨行"
)

# 添加性别判定结果统计的接口
iface3 = gr.Interface(
    count_gender_results,
    inputs="text",
    outputs="text",
    title="性别判定结果统计",
    description="根据上一步性别判定结果，输入上一步的输出进行统计"
)

# 创建从文件批量预测性别的 Gradio 接口
# 确保Gradio接口使用更新后的函数
iface4 = gr.Interface(
    predict_gender_batch_from_file,
    inputs="text",
    outputs="text",
    title="人名文件批量性别判定",
    description="输入包含人名的文件路径进行批量性别判定"
)

demo = gr.TabbedInterface(
    [iface1, iface2, iface3, iface4],
    tab_names=["单个人名性别判定", "批量人名性别判定", "批量性别判定结果统计","人名文件批量性别判定"],
    title="听其名，知其性"
)

demo.launch()